from PIL import Image
import numpy as np
import matplotlib.pyplot as plt

def get_length_info(length):
    residual_list = []
    for i in range(3):
        residual_list.append(length % 256)
        length = length // 256
    assert len(residual_list) <= 3,'length too big'
    residual_list.reverse()
    return residual_list
    pass


def save_image(image,py_file,patch=(200,400,500,700),interval=7):
    image_obj = Image.open(image)
    image_arr = np.array(image_obj)
    fid = open(py_file)
    py_ascii = [ord(x) for x in fid.read()]
    fid.close()
    py_arr = np.array(py_ascii) + 128
    py_arr_len = len(py_arr)
    length_info_list = get_length_info(py_arr_len)
    info_all = length_info_list + py_arr.tolist()
    # the first three pixels is the length of the py_arr
    # of 255 exp
    image_patch = image_arr[patch[0]:patch[1],patch[2]:patch[3],:]
    image_patch_reshape = image_patch.reshape(-1)
    image_patch_reshape[:len(info_all) * interval:interval] = info_all
    image_patch = image_patch_reshape.reshape(image_patch.shape)
    plt.imshow(image_patch)
    plt.show()
    image_arr[patch[0]:patch[1],patch[2]:patch[3],:] = image_patch
    plt.imshow(image_arr)
    plt.show()
    image_obj = Image.fromarray(image_arr)
    save_name = image.split('.')[0] + '1.jpg'
    image_obj.save(save_name,format='BMP')

def load_image(image,patch=(200,400,500,700),interval=7,save_name='test1.py'):
    image_obj = Image.open(image)
    image_arr = np.array(image_obj)
    image_patch = image_arr[patch[0]:patch[1], patch[2]:patch[3], :]
    image_patch_reshape = image_patch.reshape(-1)
    length = image_patch_reshape[0 * interval] * 256 ** 2 + image_patch_reshape[1 * interval] * 256 + \
        image_patch_reshape[2 * interval]
    image_info = image_patch_reshape[3 * interval:(3 + length) * interval:interval] - 128
    content = ''.join([chr(x) for x in image_info])
    fid = open(save_name,'w')
    fid.write(content)
    fid.close()

    pass

def generate_curves(advantage=0.001,save_path='./test.jpg'):
    plt.figure(figsize=(28,12))
    ret = np.random.rand(1000,500) / 10 - 0.05 + advantage
    cum_ret = np.cumsum(ret,axis=1)

    plt.plot(cum_ret.transpose())
    # plt.plot(ret,'+')
    # plt.show()
    plt.grid(True)
    plt.savefig(save_path)
    # plt.show()
    # plt.hist(ret,bins=1000)
    pass


def check_image(image='./test.jpg'):
    while True:
        image_obj = Image.open(image)
        image_arr = np.array(image_obj)
        input_list = [int(x) for x in input().split(' ')]
        image_arr[input_list[0]:input_list[1],input_list[2]:input_list[3],:] = 0
        plt.imshow(image_arr)
        plt.show()

        # plt.imshow(image_obj)


# curves_0 550 750 500 2500
# curves_1 200 950 400 2500



if __name__ == '__main__':
    # for i in range(1):
    #     generate_curves(advantage=(i + 1) * 0.001,save_path='curves_%d.jpg'%i)
    # save_image('./curves.jpg','Manifold.py')
    # check_image('./curves_0.jpg')
    patch = (550,750,500,2500)
    # save_image('curves_0.jpg','Manifold.py',patch=patch,interval=7)
    load_image('curves_01.jpg',patch=patch,interval=9,save_name='test1.py')
    pass
